package org.yagnus.stats.sampler.discrete.withoutreplacement;

import java.util.ArrayList;

/**
 * 
 * @author hc.busy
 * 
 *         This is a most obvious implementation using arralists and O(n) get/remove methods, but it
 *         should be correct.
 * 
 * 
 * 
 * 
 *         For large samples TODO: Check counts after repeated sample from unequal
 *         probabilities--check against brute force computed expected value. TODO: Check the order
 *         in which an item is drawn after a large number of takeAll by storing the # of draw rank.
 *         (compute draw rank theoretcially and compare)
 * 
 */
public class ListArraySampler<T> extends WithoutReplacementSampler<T> {

    public ListArraySampler(T[] t, double[] weights) {
        super(t, weights);
    }

    public ListArraySampler(T[] t, long[] weights) {
        super(t, weights);
    }

    public ListArraySampler(T[] t, int[] weights) {
        super(t, weights);
    }

    public ListArraySampler(T[] t) {
        super(t);
    }

    /**
     * 
     * @returns next sample or null on error.
     */
    @Override
    public T take() {
        if (samps.size() == 0 || curTotal == 0) {
            return null;
        }
        T ret = null;
        double myTotal = 0;
        do {

            for (int i = 0; i < samps.size(); ++i) {
                double d = rng().nextDouble() * curTotal;

                myTotal += p.get(i);
                if (myTotal >= d) {
                    // got our sample.
                    ret = actualSamps.remove(i);
                    curTotal -= p.remove(i);
                    break;
                }
            }
        } while (ret == null);
        return ret;
    }
    ArrayList<Double> p;
    ArrayList<T> samps;
    ArrayList<Double> actualP;
    ArrayList<T> actualSamps;
    double curTotal = 0;

    @Override
    protected void _init(T[] t, double[] weights) {

        if (t.length != weights.length) {
            throw new IllegalArgumentException("weights not same length as sample space.");
        }
        for (int i = 0; i < t.length; ++i) {
            T tv = t[i];
            if (tv == null) {
                throw new IllegalArgumentException("null weights are not suported.");
            }
            this.p.add(weights[i]);
            this.samps.add(t[i]);
        }
        _reset();

    }

    protected void _reset() {
        curTotal = 0;
        actualP = new ArrayList<Double>(p.size());
        actualSamps = new ArrayList<T>(samps.size());

        for (Double d : p) {
            double v = d.doubleValue();
            curTotal += v;
            actualP.add(v);
        }

        for (T t : samps) {
            actualSamps.add(t);
        }
    }

    @Override
    public void addSample(T t, double w) {
        samps.add(t);
        p.add(w);
        actualSamps.add(t);
        actualP.add(w);
        curTotal += w;
    }

    @Override
    public void removeSample(T t) {

        // Remove from original samples pool.
        for (int i = 0; i < p.size(); ++i) {
            if (samps.get(i).equals(t)) {
                samps.remove(i);
                p.remove(i);
                --i;
            }
        }

        // remove from current remaining sample pool
        for (int i = 0; i < actualP.size(); ++i) {
            if (samps.get(i).equals(t)) {
                samps.remove(i);
                curTotal -= p.remove(i);
                --i; // keep i same since we deleted last one.
            } // now remove from part of the list outside of the available sample.
        }
    }
}
